# -*- coding: utf-8 -*-
"""
Created on Tue Jun 16 06:29:09 2020

@author: HP
"""


from os.path import isfile
import numpy as np
import pandas as pd
from itertools import permutations
from numpy.random import default_rng
from increasing_closures import lottery_equivalence, \
    impatience, correlation_neutrality, correlation_aversion, \
    correlation_seeking, stochastic_impatience, stochastic_patience
from garp_NOQ import garp_NOQ_find_e

#=============================================================================
#VALIDATION FUNCTIONS
#=============================================================================

def impatience_ok( q_vec, p_vec ):
    result = True
    if (p_vec[0] <= p_vec[1]) & (q_vec[0] < q_vec[1]):
        result = False
    elif (p_vec[2] <= p_vec[3]) & (q_vec[2] < q_vec[3]):
        result = False
    return result
    

def fosd_ok( q_vec, p_vec ):
    result = True
    q_mirror_vec = np.append( q_vec[[2,3]], q_vec[[0,1]] )
    if p_vec @ q_vec > p_vec @ q_mirror_vec:
        result = False
    return result

#=============================================================================
#SIMULATE FILE
#=============================================================================
sim_file = 'simulate_edge.xlsx'
simN = 10000

if isfile(sim_file):
    sim_df = pd.read_excel(sim_file, sheet_name = 0, index_col = 0)
else:
    cols = ['participantid', 'x', 'y', 'z', 'w', 'px', 'py', 'pz', 'pw']
    #sim_df = pd.DataFrame(np.nan, columns = cols, index = range(41*simN))
    #sim_df['participantid'] = np.floor( sim_df.index / 41. ).astype('int')+1  
    #==========================================================================
    #MAKE PRICE LIST
    #==========================================================================
    
    #Make the 96 price vectors
    price_list = [0.5,0.8,1.25,2.]
    prices_df = pd.DataFrame( np.nan, columns = ['px','py', 'pz','pw'], index = range(96) )
    rowi = 0 #keep track of the row I'm working on
    for ones_place in range(4):
        #ones_place marks where the 1 will appear.
        #price_list_active holds different 3-permutations of the price_list.
        for price_list_active in permutations( price_list, 3 ):
            prices_df.iloc[ rowi, ones_place ] = 1.
            prices_df.iloc[ rowi, np.array([0,1,2,3]) != ones_place ] = np.array( price_list_active )
            rowi = rowi + 1
    
    #==========================================================================
    #MAKE THE PRICES FOR SIM_DF
    #==========================================================================
    sim_df_list = []
    for simi in range(simN):
        asim_df = pd.DataFrame( np.nan, columns = cols, index = range(41 * simi, 41 * simi + 41) )
        asim_df['participantid'] = simi + 1
        asim_df.iloc[0, 5:9] = np.array([1.,1.,1.,1.])
        asim_df.iloc[1:41, 5:9] = prices_df.sample(n=40, replace=False).values
        sim_df_list.append(asim_df)
    sim_df = pd.concat( sim_df_list )
    sim_df.to_excel(sim_file)

#==========================================================================
#GENERATE CHOICE DATA
#==========================================================================

rowN = sim_df.shape[0]
rng = default_rng()

changedFlag = False
for rowi in range(rowN):
    if pd.isnull(sim_df.iloc[rowi,1:5]).any():
        changedFlag = True
        #If the value is missing then fill it
        p_vec = sim_df.iloc[rowi,5:9].values
        flag = False #Turn to true once choices satisfy FOSD and patience
        while not flag:
            b_vec = rng.dirichlet( [1., 1.] )
            non_zeros = rng.choice([0,1,2,3],2,replace=False)
            q_vec = np.zeros(4)
            q_vec[non_zeros[0]] = 100. * ( b_vec[0] / p_vec[ non_zeros[0] ] ) 
            q_vec[non_zeros[1]] = 100. * ( b_vec[1] / p_vec[ non_zeros[1] ] )            
            if impatience_ok(q_vec = q_vec, p_vec = p_vec) & fosd_ok(q_vec = q_vec, p_vec = p_vec):
                flag = True
        sim_df.iloc[rowi,1:5] = q_vec
    if not ( (rowi + 1) % 100 ):
        print('rowi: {0}'.format(rowi))
        
    if changedFlag & ( not ( (rowi+1) % 1000)):
        print('rowi: {0}'.format(rowi))
        sim_df.to_excel(sim_file)
if changedFlag:
    sim_df.to_excel(sim_file)
    
#==========================================================================
#RESULTS FILE
#==========================================================================

imple_list = [impatience, lottery_equivalence]
tests_dict = {'def': [], \
              'cn': [correlation_neutrality], \
              'ca': [correlation_aversion], \
              'cs': [correlation_seeking], \
              'si': [stochastic_impatience], \
              'sp': [stochastic_patience]}

columns = list(tests_dict.keys())
pids = sim_df['participantid'].unique()

results_file = 'rp_tests_by_participant_edge_sim.xlsx'

#-------------------
#manage results file
#-------------------
if isfile( results_file ):
    results_df = pd.read_excel( results_file, sheet_name = 0, index_col = 0 )
else:
    #Make a new results df
    results_df = pd.DataFrame( np.nan, index = pids, columns = columns )
    results_df.index.name = 'participantid'

#==========================================================================
#RUN TESTS
#==========================================================================
    
counter = 0
for pid in pids:
    if pd.isnull( results_df.loc[pid,:] ).any():
        print('current id: {0}'.format(pid))
        counter = counter + 1
        cur_df = sim_df[ sim_df['participantid'] == pid ]
        C = cur_df[['x','y','z','w']].to_numpy()
        P = cur_df[['px','py','pz','pw']].to_numpy()
        for test_name in tests_dict:
            #Run the test if we don't already have the results.
            e = garp_NOQ_find_e(C,P,Ilist= tests_dict[test_name])
            results_df.loc[pid,test_name] = e
        if counter == 100:
            print('saving results')
            results_df.to_excel(results_file)
            counter = 0
if counter:
    results_df.to_excel(results_file)
    counter = 0

#==========================================================================
#SUMMARIZE RESULTS
#==========================================================================

summary_file = 'simulate_edge_summary.xlsx'

#columns and row names
escores = [100,99,95,90]

var_names = list(tests_dict.keys())

summary_df = pd.DataFrame(data=np.nan, index = escores, columns = var_names)

for escore in escores:

    elevel_df = results_df.copy()
    elevel_df[elevel_df < (escore / 100.)] = 0
    elevel_df[elevel_df >= (escore / 100.)] = 1
    summary_df.loc[ escore, : ] = elevel_df.mean()
    
summary_df.to_excel(summary_file)




